package com.gitee.feizns.explore.data_structure.tree.cus.bst;

import com.gitee.feizns.explore.data_structure.tree.cus.Node;
import com.gitee.feizns.explore.data_structure.tree.cus.PutRet;
import com.gitee.feizns.explore.data_structure.tree.cus.TreeUtils;

import java.util.Collection;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;

/**
 * @author feizns
 * @since 2019/11/25
 */
public class BstImpl<K extends Comparable<K>, V> implements Bst<BstImpl.NodeImpl, K, V> {

    private NodeImpl<K, V> root;

    private int size;

    @Override
    public int size() {
        return size;
    }

    @Override
    public V get(Object key) {
        if ( key != null && key instanceof Comparable ) {
            NodeImpl<K, V> kvNode = TreeUtils.find(root, (K) key);
            return kvNode != null ? kvNode.val() : null;
        }
        return null;
    }

    @Override
    public V put(K key, V value) {
        if ( key == null )
            throw new NullPointerException("Key");
        PutRet<NodeImpl, V> ret = put(root, key, value);
        root = ret.node();
        return ret.v();
    }

    private PutRet<NodeImpl, V> put(NodeImpl node, K key, V value) {
        PutRet<NodeImpl, V> ret = null;
        if ( node != null ) {
            int weight = node.key().compareTo(key);

            if ( weight == 0 ) {
                V v = (V) node.val();
                node.val(value);
                ret = new PutRet<>(node, v);
            } else if ( weight > 0 ) {
                PutRet<NodeImpl, V> tmp = put(node.left(), key, value);
                node.left(tmp.node());
                ret = new PutRet<>(node, tmp.v());
            } else {
                PutRet<NodeImpl, V> tmp = put(node.right(), key, value);
                node.right(tmp.node());
                ret = new PutRet<>(node, tmp.v());
            }
        }

        if ( ret == null ) {
            size++;
            return new PutRet<>(new NodeImpl<>(key, value), null);
        }
        return ret;
    }

    @Override
    public Node root() {
        return root;
    }

    @Override
    public V remove(Object key) {
        if ( key != null && key instanceof Comparable ) {
            PutRet<NodeImpl, V> putRet = TreeUtils.remove(root, (K) key);
            root = putRet.node();
            return putRet.v();
        }
        return null;
    }

    @Override
    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public Set<K> keySet() {
        return entrySet().stream().map(Entry::getKey).collect(Collectors.toCollection(TreeSet::new));
    }

    @Override
    public Collection<V> values() {
        return entrySet().stream().map(Entry::getValue).collect(Collectors.toList());
    }

    @Override
    public Set<Entry<K, V>> entrySet() {
        Set<Entry<K, V>> entries = new TreeSet<>();
        recursiveMiddleEach((k, v) -> entries.add(new EntryImpl<>(k, v)));
        return entries;
    }

    static class NodeImpl<K extends Comparable<K>, V> implements Node<NodeImpl, K, V> {

        K k;

        V v;

        NodeImpl<K, V> left, right;

        public NodeImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }

        @Override
        public K key() {
            return k;
        }

        @Override
        public void key(K k) {
            this.k = k;
        }

        @Override
        public V val() {
            return v;
        }

        @Override
        public void val(V v) {
            this.v = v;
        }

        @Override
        public NodeImpl<K, V> left() {
            return left;
        }

        @Override
        public void left(NodeImpl left) {
            this.left = left;
        }

        @Override
        public NodeImpl<K, V> right() {
            return right;
        }

        @Override
        public void right(NodeImpl right) {
            this.right = right;
        }

    }

    private static class EntryImpl<K extends Comparable<K>, V> implements Entry<K, V> {

        K k;

        V v;

        @Override
        public K getKey() {
            return k;
        }

        @Override
        public V getValue() {
            return v;
        }

        @Override
        public V setValue(V value) {
            V old = this.v;
            this.v = value;
            return old;
        }

        public EntryImpl(K k, V v) {
            this.k = k;
            this.v = v;
        }

    }

}
